

def mask_select(inputs, mask):
    input_dim = inputs.ndim
    mask_dim = mask.ndim
    mask = mask.flatten().astype('bool')
    if input_dim > mask_dim:
        inputs = inputs.reshape([*mask.shape, -1])[mask]
    else:
        inputs = inputs.flatten()[mask]
    return inputs


def mask_loss(func):
    def new_func(self, inputs, targets, mask):
        if mask is not None:
            inputs = mask_select(inputs, mask)
            targets = mask_select(targets, mask)
        return func(self, inputs, targets)

    return new_func


def shift_right_mask_loss(func):
    def new_func(self, inputs, targets, mask):
        mask = mask[:, 1:]
        inputs = inputs[:, :-1]
        targets = targets[:, 1:]
        inputs = mask_select(inputs, mask)
        targets = mask_select(targets, mask)
        return func(self, inputs, targets)

    return new_func
